﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace GrayParrot.Extensions.Queries
{
    public static class QueriesExtensions
    {
        public static IEnumerable<T> Flatten<T>(this T source, Func<T, IEnumerable<T>> selector)
        {
            var stack = new Stack<T>();
            stack.Push(source);
            while (stack.Count > 0)
            {
                var current = stack.Pop();
                yield return current;
                foreach (var child in selector(current).Reverse())
                    stack.Push(child);
            }
        }

        public static IEnumerable<T> Descendants<T>(this IEnumerable<T> source, Func<T, IEnumerable<T>> descendBy)
        {
            if (source != null)
            {
                foreach (T value in source)
                {
                    yield return value;

                    if (descendBy(value) != null)
                    {
                        foreach (T child in descendBy(value).Descendants<T>(descendBy))
                        {
                            yield return child;
                        }
                    }
                }
            }
        }
    }
}
